#!/usr/bin/env python3
import sys
import os
import subprocess
import tempfile
import re
c = False
v = False
k = False
vd = False
o = None
src = None
modulepath = None
args = []
prefix = os.path.dirname(os.path.abspath(sys.argv[0]))
plugin = os.path.join(prefix, "omp2ath.so")
lib = os.path.join(prefix, "libomp2ath.o")
O2ATH_RE = re.compile(r"\s*\!\$\s*o2ath\s+entry\((.*)\)\s*")
i = 1
argv = list(filter(lambda x: not x.startswith("-Wslave,") and not x.startswith("-Whost",), sys.argv))
args_slave = []
args_host = []
while i < len(argv):
    arg = argv[i].strip()
    #print(i, arg)
    if arg == "-c":
        c = True
    elif arg == "-v":
        v = True
    elif arg == "-vv":
        args.append("-v")
        v = True
    elif arg == "-k":
        k = True
    elif arg == "-o":
        o = argv[i+1]
        i += 1
    elif arg == "-vd":
        vd = True
    elif arg.startswith("-J"):
        modulepath = arg[2:]
    elif os.path.splitext(arg)[-1] in [".f90", ".F90", ".F", ".f", ".in"]:
        src = arg
        args.append(arg)
    else:
        args.append(arg)
    i += 1
for slf in filter(lambda x: x.startswith("-Wslave,"), sys.argv):
    args_slave += slf.split(",")[1:]
for hof in filter(lambda x: x.startswith("-Whost,"), sys.argv):
    args_host += hof.split(",")[1:]
def run(args, fatal=True):
    if v:
        print(args, file=sys.stderr)
    ret = subprocess.run(args, stderr=sys.stderr, stdout=sys.stdout).returncode
    if fatal and ret:
        print(str(args) + " returns non-zero status %d" % ret, file=sys.stderr)
        sys.exit(ret)
    return ret
def check(args):
    if v:
        print(args, file=sys.stderr)
    return subprocess.check_output(args, stderr=sys.stderr).decode()
if o is None:
    if src is not None:
        if c:
            o = os.path.splitext(os.path.basename(src))[0] + ".o"
        else:
            o = "a.out"
if modulepath is None:
    modulepath = "."
if src is None:
    sys.exit(run(["mpif90", "-fopenmp", "-Wl,--wrap,__real_athread_spawn"]  + [lib] + argv[1:]))
try:
    work = tempfile.mkdtemp(prefix="xfort")
    obase = os.path.basename(os.path.splitext(o)[0])
    host_o = os.path.join(work, obase + ".host.o")
    slave_o = os.path.join(work, obase + ".slave.o")
    slave_flag = os.path.join(work, "slave_flag")
    os.environ["SLAVE_FLAG"] = slave_flag
    entries = []
    if src is not None:
        for line in open(src):
            if O2ATH_RE.fullmatch(line):
                entries = entries + list(map(str.strip, O2ATH_RE.fullmatch(line).group(1).split(",")))
    os.environ["O2ATH_ENTRIES"] = ":".join(entries)
    if entries:
        print("making the plugin to consider %s for slave code generation" % str(entries), file=sys.stderr)
    run(["mpif90", "-mhost", "-c", "-fplugin=%s" % plugin, "-fopenmp", "-D__sw_host__", "-J%s"%work] + args + args_host + ["-o", host_o])
    if os.path.exists(slave_flag) or entries:
        run(["sw9gfortran", "-mslave", "-c", "-fplugin=%s" % plugin, "-fopenmp", "-I/usr/sw/mpi/mpi_current/include/", "-D__sw_slave__", "-J%s"%work] + args + args_slave + ["-o", slave_o])
        if c:
            run(["sw9ld", "-r", "-z", "muldefs", host_o, slave_o, "-o", o])
        else:
            run(["mpif90", "-mhybrid", "-fopenmp", host_o, slave_o, lib, "-o", o, "-Wl,--wrap,__real_athread_spawn"])
    else:
        if c:
            run(["mv", host_o, o])
        else:
            run(["mpif90", "-mhybrid", host_o, "-o", o])
    for f in os.listdir(work):
        if f.endswith(".mod"):
            if not os.path.exists(os.path.join(modulepath, f)) or open(os.path.join(work, f), "rb").read() != open(os.path.join(modulepath, f), "rb").read():
                ov = v
                v = True
                run(["cp", "-f", os.path.join(work, f), os.path.join(modulepath, f)])
                v = ov
finally:
    if not k:
        run (["rm", "-rf", work])
    else:
        print("temp files saved in %s" % work)
#if not k:
    #run(["rm"] + temps + ["-v"] if v else [])
